#!/usr/bin/env python3
"""
支持模板化配置，便于维护多个机箱
"""

import json
import time
import argparse
import numpy as np
from typing import Dict, List, Optional, Any
import epics


class SimpleBMPProcessor:

    
    def __init__(self, config_file: str):
        """初始化处理器"""
        self.config_file = config_file
        self.config = self._load_config()
        self.pvs = {}
        self.connected = False
        
        # 连接PV
        if not self.connect_pvs():
            raise RuntimeError("无法连接到所有必需的PV")
        
    def _load_config(self) -> Dict[str, Any]:
        """加载配置文件"""
        try:
            with open(self.config_file, 'r', encoding='utf-8') as f:
                config = json.load(f)
            
            # 检查是否是模板化配置
            if 'chassis_list' in config and 'pv_templates' in config:
                return self._parse_template_config(config)
            else:
                # 兼容旧的配置格式
                return config
                
        except FileNotFoundError:
            raise FileNotFoundError(f"配置文件 {self.config_file} 不存在")
        except json.JSONDecodeError as e:
            raise ValueError(f"配置文件格式错误: {e}")
    
    def _parse_template_config(self, template_config: Dict[str, Any]) -> Dict[str, Any]:
        """解析模板化配置。
        
        此方法根据模板配置和机箱列表生成最终的配置字典。
        它会为每个BPM对生成波形PV、平均值PV和位置设置PV的地址。
        """
        config = {
            'slice_start': template_config.get('slice_start', 0),
            'slice_end': template_config.get('slice_end', 100),
            'connection_timeout': template_config.get('connection_timeout', 5.0),
            'chassis': {}
        }
        
        # 获取PV模板，这些模板定义了如何根据前缀和BPM名称生成PV地址
        templates = template_config['pv_templates']
        
        # 遍历配置文件中定义的每个机箱
        for chassis_info in template_config['chassis_list']:
            chassis_name = chassis_info['name']
            prefix = chassis_info['prefix']
            bmp_pairs = chassis_info['bmp_pairs'] # 每个机箱可能包含多个BPM对
            
            chassis_config = {'bmps': {}}
            
            # 遍历机箱中的每个BPM对
            for bmp_pair in bmp_pairs:
                bmp1_name = bmp_pair['bmp1']
                bmp2_name = bmp_pair['bmp2']
                
                # 为BPM1生成所有相关的PV地址
                chassis_config['bmps'][bmp1_name] = {
                    'waveforms': self._substitute_template(
                        templates['bmp1_waveforms'], 
                        prefix=prefix, 
                        bmp1=bmp1_name, 
                        bmp2=bmp2_name
                    ), # 波形数据PV
                    'averages': self._substitute_template(
                        templates['bmp1_averages'], 
                        prefix=prefix, 
                        bmp1=bmp1_name, 
                        bmp2=bmp2_name
                    ), # 平均值输出PV
                    'positions': self._substitute_template(
                        templates['bmp1_positions'], 
                        prefix=prefix, 
                        bmp1=bmp1_name, 
                        bmp2=bmp2_name.replace("BPM", "")
                    ) # 有效信号取值区间设置PV
                }
                
                # 为BPM2生成所有相关的PV地址
                chassis_config['bmps'][bmp2_name] = {
                    'waveforms': self._substitute_template(
                        templates['bmp2_waveforms'], 
                        prefix=prefix, 
                        bmp1=bmp1_name, 
                        bmp2=bmp2_name
                    ), # 波形数据PV
                    'averages': self._substitute_template(
                        templates['bmp2_averages'], 
                        prefix=prefix, 
                        bmp1=bmp1_name, 
                        bmp2=bmp2_name
                    ), # 平均值输出PV
                    'positions': self._substitute_template(
                        templates['bmp2_positions'], 
                        prefix=prefix, 
                        bmp1=bmp1_name, 
                        bmp2=bmp2_name.replace("BPM", "")
                    ) # 有效信号取值区间设置PV
                }
            
            config['chassis'][chassis_name] = chassis_config
        
        return config
    
    def _substitute_template(self, template: Dict[str, str], **kwargs) -> Dict[str, str]:
        """替换模板中的占位符"""
        result = {}
        for key, value in template.items():
            result[key] = value.format(**kwargs)
        return result
    
    def connect_pvs(self) -> bool:
        """连接所有在配置中定义的EPICS PV。
        
        此方法遍历所有机箱和BPM，尝试连接波形PV、平均值PV和位置设置PV。
        如果任何PV连接失败，则返回False并打印错误信息。
        """
        print("正在连接PV...")
        timeout = self.config.get('connection_timeout', 5.0)
        
        for chassis_name, chassis_config in self.config['chassis'].items():
            print(f"连接机箱 {chassis_name}...")
            
            for bmp_name, bmp_config in chassis_config['bmps'].items():
                # 连接波形数据PV (waveform PVs)
                for pv_name, pv_addr in bmp_config['waveforms'].items():
                    try:
                        pv = epics.PV(pv_addr, connection_timeout=timeout)
                        if pv.wait_for_connection(timeout=timeout):
                            self.pvs[f"{chassis_name}_{bmp_name}_{pv_name}"] = pv
                            print(f"  ✓ {pv_addr}")
                        else:
                            print(f"  ✗ {pv_addr} - 连接超时")
                            return False
                    except Exception as e:
                        print(f"  ✗ {pv_addr} - 连接错误: {e}")
                        return False
                
                # 连接平均值输出PV (average PVs)
                for pv_name, pv_addr in bmp_config['averages'].items():
                    try:
                        pv = epics.PV(pv_addr, connection_timeout=timeout)
                        if pv.wait_for_connection(timeout=timeout):
                            self.pvs[f"{chassis_name}_{bmp_name}_{pv_name}"] = pv
                            print(f"  ✓ {pv_addr}")
                        else:
                            print(f"  ✗ {pv_addr} - 连接超时")
                            return False
                    except Exception as e:
                        print(f"  ✗ {pv_addr} - 连接错误: {e}")
                        return False
                
                # 连接位置设置PV (position PVs)
                for pv_name, pv_addr in bmp_config['positions'].items():
                    try:
                        pv = epics.PV(pv_addr, connection_timeout=timeout)
                        if pv.wait_for_connection(timeout=timeout):
                            self.pvs[f"{chassis_name}_{bmp_name}_{pv_name}"] = pv
                            print(f"  ✓ {pv_addr}")
                        else:
                            print(f"  ✗ {pv_addr} - 连接超时")
                            return False
                    except Exception as e:
                        print(f"  ✗ {pv_addr} - 连接错误: {e}")
                        return False
        
        self.connected = True
        print("所有PV连接成功!")
        return True
    
    def calculate_slice_average(self, waveform_data: np.ndarray, start_pos: int = None, end_pos: int = None) -> float:
        """计算给定波形数据的切片平均值。
        
        如果提供了start_pos和end_pos，则使用它们作为切片范围。
        否则，使用配置文件中定义的默认slice_start和slice_end。
        
        Args:
            waveform_data (np.ndarray): 输入的波形数据数组。
            start_pos (int, optional): 切片起始位置（包含）。默认为None。
            end_pos (int, optional): 切片结束位置（不包含）。默认为None。
            
        Returns:
            float: 计算出的切片平均值。如果数据为空或切片范围无效，则返回0.0。
        """
        if waveform_data is None:
            return 0.0
        
        # 确保输入数据是numpy数组
        if not isinstance(waveform_data, np.ndarray):
            try:
                waveform_data = np.array(waveform_data)
            except Exception:
                return 0.0
        
        # 检查数组是否为空或者是标量
        if waveform_data.size == 0:
            return 0.0
        
        # 如果是标量，转换为单元素数组
        if waveform_data.ndim == 0:
            waveform_data = np.array([waveform_data])
        
        if len(waveform_data) == 0:
            return 0.0
        
        # 如果提供了具体的位置参数，使用它们；否则使用配置文件中的默认值
        if start_pos is not None and end_pos is not None:
            start = max(0, start_pos)
            end = min(len(waveform_data), end_pos)
        else:
            start = max(0, self.config['slice_start'])
            end = min(len(waveform_data), self.config['slice_end'])
        
        if start >= end:
            return 0.0
        
        slice_data = waveform_data[start:end]
        return float(np.mean(slice_data))
    
    def process_single(self):
        """执行单次BMP数据处理，包括背景减除和平均值计算。
        
        此方法遍历所有配置的BMP，读取其波形数据，并根据配置的切片范围和
        从PV读取的有效信号范围计算背景减除后的BMP电压。
        计算结果将写入对应的平均值PV。
        
        Returns:
            Dict[str, Dict[str, float]]: 包含每个机箱和BMP处理结果的字典。
        """
        if not self.connected:
            print("错误: PV未连接。请先调用connect_pvs()。")
            return {}

        results = {}
        
        for chassis_name, chassis_config in self.config['chassis'].items():
            results[chassis_name] = {}
            
            for bmp_name, bmp_config in chassis_config['bmps'].items():
                print(f"处理 {chassis_name}/{bmp_name}...")
                
                try:
                    # 从PV读取有效信号范围
                    avg_start_pos_pv = self.pvs[f"{chassis_name}_{bmp_name}_avg_start"].get()
                    avg_stop_pos_pv = self.pvs[f"{chassis_name}_{bmp_name}_avg_stop"].get()

                    effective_start_pos = None
                    effective_stop_pos = None

                    if avg_start_pos_pv is not None and avg_stop_pos_pv is not None:
                        try:
                            start_int = int(avg_start_pos_pv)
                            stop_int = int(avg_stop_pos_pv)
                            if start_int < stop_int:
                                effective_start_pos = start_int
                                effective_stop_pos = stop_int
                            else:
                                print(f"警告: {chassis_name}/{bmp_name} 的有效信号区间设置无效 ({start_int}-{stop_int})，将使用默认值 {self.config['slice_start']}-{self.config['slice_end']}。")
                        except ValueError:
                            print(f"警告: {chassis_name}/{bmp_name} 的有效信号区间PV值无法转换为整数，将使用默认值 {self.config['slice_start']}-{self.config['slice_end']}。")
                    else:
                        print(f"警告: {chassis_name}/{bmp_name} 的有效信号区间PV值未设置，将使用默认值 {self.config['slice_start']}-{self.config['slice_end']}。")
                    
                    # 获取波形数据 (确保连接稳定后再获取数据)
                    # 等待一小段时间确保PV连接完全稳定
                    time.sleep(0.1)
                    
                    waveform_va = self.pvs[f"{chassis_name}_{bmp_name}_va"].get(use_monitor=False, timeout=2.0)
                    waveform_vb = self.pvs[f"{chassis_name}_{bmp_name}_vb"].get(use_monitor=False, timeout=2.0)
                    waveform_vc = self.pvs[f"{chassis_name}_{bmp_name}_vc"].get(use_monitor=False, timeout=2.0)
                    waveform_vd = self.pvs[f"{chassis_name}_{bmp_name}_vd"].get(use_monitor=False, timeout=2.0)
                    
                    if waveform_va is None or waveform_vb is None or waveform_vc is None or waveform_vd is None:
                        print(f"警告: 无法获取 {chassis_name}/{bmp_name} 的波形数据。跳过。")
                        results[chassis_name][bmp_name] = {"va_average": 0.0, "vb_average": 0.0, "vc_average": 0.0, "vd_average": 0.0}
                        continue
                    
                    # 计算背景平均值 (使用配置文件中的背景段范围)
                    background_va = self.calculate_slice_average(waveform_va, self.config['slice_start'], self.config['slice_end'])
                    background_vb = self.calculate_slice_average(waveform_vb, self.config['slice_start'], self.config['slice_end'])
                    background_vc = self.calculate_slice_average(waveform_vc, self.config['slice_start'], self.config['slice_end'])
                    background_vd = self.calculate_slice_average(waveform_vd, self.config['slice_start'], self.config['slice_end'])
                    
                    # 打印背景段和有效信号段信息
                    print(f"  背景段范围: {self.config['slice_start']} - {self.config['slice_end']}")
                    if effective_start_pos is not None and effective_stop_pos is not None:
                        print(f"  有效信号段: {effective_start_pos} - {effective_stop_pos}")
                    else:
                        print(f"  有效信号段: {self.config['slice_start']} - {self.config['slice_end']} (默认值)")
                        # 如果没有有效的PV设置，使用默认值但确保与背景段不同
                        effective_start_pos = 5  # 使用固定的信号段范围
                        effective_stop_pos = 15
                        print(f"  警告: 使用固定信号段范围 {effective_start_pos} - {effective_stop_pos}")
                    
                    # 计算信号平均值 (使用从PV读取的有效信号范围)
                    signal_va = self.calculate_slice_average(waveform_va, effective_start_pos, effective_stop_pos)
                    signal_vb = self.calculate_slice_average(waveform_vb, effective_start_pos, effective_stop_pos)
                    signal_vc = self.calculate_slice_average(waveform_vc, effective_start_pos, effective_stop_pos)
                    signal_vd = self.calculate_slice_average(waveform_vd, effective_start_pos, effective_stop_pos)
                    
                    # 打印信号段的原始数据
                    if waveform_va is not None:
                        # 检查波形数据类型和长度
                        print(f"  波形数据类型: Va={type(waveform_va)}, Vb={type(waveform_vb)}, Vc={type(waveform_vc)}, Vd={type(waveform_vd)}")
                        
                        # 如果是数组类型，打印信号段数据
                        if hasattr(waveform_va, '__len__') and not isinstance(waveform_va, str):
                            if len(waveform_va) > effective_stop_pos:
                                signal_data_va = waveform_va[effective_start_pos:effective_stop_pos]
                                signal_data_vb = waveform_vb[effective_start_pos:effective_stop_pos]
                                signal_data_vc = waveform_vc[effective_start_pos:effective_stop_pos]
                                signal_data_vd = waveform_vd[effective_start_pos:effective_stop_pos]
                                
                                print(f"  信号段数据 Va[{effective_start_pos}:{effective_stop_pos}]: {signal_data_va}")
                                print(f"  信号段数据 Vb[{effective_start_pos}:{effective_stop_pos}]: {signal_data_vb}")
                                print(f"  信号段数据 Vc[{effective_start_pos}:{effective_stop_pos}]: {signal_data_vc}")
                                print(f"  信号段数据 Vd[{effective_start_pos}:{effective_stop_pos}]: {signal_data_vd}")
                            else:
                                print(f"  警告: 波形数据长度不足，无法提取信号段 [{effective_start_pos}:{effective_stop_pos}]")
                                print(f"  波形数据长度: Va={len(waveform_va)}, Vb={len(waveform_vb)}, Vc={len(waveform_vc)}, Vd={len(waveform_vd)}")
                        else:
                            # 如果是标量值，直接打印
                            print(f"  波形数据为标量值: Va={waveform_va}, Vb={waveform_vb}, Vc={waveform_vc}, Vd={waveform_vd}")
                    else:
                        print(f"  警告: 波形数据为None")
                    
                    # 打印有效信号平均值和背景平均值
                    print(f"  有效信号平均值 - Va: {signal_va:.6f}, Vb: {signal_vb:.6f}, Vc: {signal_vc:.6f}, Vd: {signal_vd:.6f}")
                    print(f"  背景平均值 - Va: {background_va:.6f}, Vb: {background_vb:.6f}, Vc: {background_vc:.6f}, Vd: {background_vd:.6f}")
                    
                    # 计算背景减除后的BMP电压
                    bmp_voltage_va = signal_va - background_va
                    bmp_voltage_vb = signal_vb - background_vb
                    bmp_voltage_vc = signal_vc - background_vc
                    bmp_voltage_vd = signal_vd - background_vd
                    
                    # 将结果写入平均值PV
                    self.pvs[f"{chassis_name}_{bmp_name}_va"].put(bmp_voltage_va)
                    self.pvs[f"{chassis_name}_{bmp_name}_vb"].put(bmp_voltage_vb)
                    self.pvs[f"{chassis_name}_{bmp_name}_vc"].put(bmp_voltage_vc)
                    self.pvs[f"{chassis_name}_{bmp_name}_vd"].put(bmp_voltage_vd)
                    
                    print(f"  Va电压: {bmp_voltage_va:.4f}, Vb电压: {bmp_voltage_vb:.4f}, Vc电压: {bmp_voltage_vc:.4f}, Vd电压: {bmp_voltage_vd:.4f}")
                    
                    # 保存结果
                    results[chassis_name][bmp_name] = {
                        "va_average": bmp_voltage_va,
                        "vb_average": bmp_voltage_vb,
                        "vc_average": bmp_voltage_vc,
                        "vd_average": bmp_voltage_vd
                    }
                    
                except Exception as e:
                    print(f"处理 {chassis_name}/{bmp_name} 时发生错误: {e}")
                    import traceback
                    traceback.print_exc()
                    results[chassis_name][bmp_name] = {"x_average": 0.0, "y_average": 0.0}
        
        return results



def main():
    """主函数，用于解析命令行参数，初始化BMP处理器，并启动连续处理循环。
    
    支持的命令行参数:
    --config: 指定配置文件的路径，默认为 'config.json'。
    --interval: 指定处理循环的间隔时间（秒），默认为1秒。
    """
    parser = argparse.ArgumentParser(description='BMP电压处理')
    parser.add_argument('--config', '-c', default='config.json',
                       help='配置文件路径 (默认: config.json)')
    parser.add_argument('--interval', '-i', type=float, default=1.0,
                       help='处理间隔(秒)')
    
    args = parser.parse_args()
    
    try:
        # 创建处理器（构造函数中已经连接PV）
        processor = SimpleBMPProcessor(args.config)
        
        print(f"开始连续处理，间隔: {args.interval}秒")
        print("按 Ctrl+C 停止...")
        
        try:
            while True:
                start_time = time.time()
                results = processor.process_single()
                
                # 打印结果摘要
                total_bmps = sum(len(chassis_results) for chassis_results in results.values())
                print(f"处理完成 - {len(results)}个机箱, {total_bmps}个BMP")
                
                # 等待下一次处理
                elapsed = time.time() - start_time
                sleep_time = max(0, args.interval - elapsed)
                if sleep_time > 0:
                    time.sleep(sleep_time)
                    
        except KeyboardInterrupt:
            print("\n停止连续处理")
            
    except Exception as e:
        print(f"错误: {e}")


if __name__ == "__main__":
    main()